import cv2
import numpy as np
from PIL import Image
import pytesseract
import logging

# 设置日志配置
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def preprocess_image(image_path):
    """对图像进行预处理（灰度化和二值化）"""
    try:
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图像文件: {image_path}")

        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        _, binary = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY_INV)
        return binary
    except Exception as e:
        logging.error(f"图像预处理失败: {e}")
        raise


def get_text_lines(binary_image):
    """获取文本行的边界框"""
    try:
        contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        bounding_boxes = []

        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            if h > 10:  # 过滤掉高度小于一定阈值的小区域
                bounding_boxes.append((x, y, w, h))

        # 按y坐标排序，确保从上到下读取
        bounding_boxes.sort(key=lambda b: b[1])
        return bounding_boxes
    except Exception as e:
        logging.error(f"获取文本行失败: {e}")
        raise


def ocr_line(line_image, lang='chi_sim'):
    """对单行图像进行OCR识别"""
    try:
        custom_config = r'--oem 3 --psm 6 -l {}'.format(lang)
        text = pytesseract.image_to_string(Image.fromarray(line_image), config=custom_config)
        return text.strip()
    except Exception as e:
        logging.error(f"单行OCR识别失败: {e}")
        raise


def ocr_core(image_path, lang='chi_sim'):
    """主函数：逐行提取图片中的文本"""
    try:
        binary_image = preprocess_image(image_path)
        bounding_boxes = get_text_lines(binary_image)

        lines = []
        for box in bounding_boxes:
            x, y, w, h = box
            line_image = binary_image[y:y + h, x:x + w]
            text = ocr_line(line_image, lang)
            lines.append(text)

        return lines
    except Exception as e:
        logging.error(f"OCR核心处理失败: {e}")
        raise


def jpg_to_png(input_jpg_path, output_png_path):
    """
    将JPG格式的图片转换为PNG格式。

    :param input_jpg_path: 输入的JPG文件路径
    :param output_png_path: 输出的PNG文件路径
    """
    try:
        # 打开JPG图像
        with Image.open(input_jpg_path) as im:
            # 保存为PNG格式
            im.save(output_png_path, 'PNG')
            print(f"成功将 {input_jpg_path} 转换为 {output_png_path}")
    except Exception as e:
        print(f"转换失败: {e}")

if __name__ == "__main__":
    # 示例调用
    try:
        # image_jpg_path = '/home/weiqiangren/Python/project/py-basic-exercises/图片文字识别/17ba29a8755a4ecbc2f698ccadf5030e_compress.jpg'  # 替换为你的图片路径
        image_png_path = '/home/weiqiangren/Python/project/py-basic-exercises/图片文字识别/企业微信截图_17387351183679.png'  # 替换为你的图片路径
        # jpg_to_png(image_jpg_path, image_png_path)
        result = ocr_core(image_png_path)
        for i, line in enumerate(result):
            print(f"Line {i + 1}: {line}")
    except Exception as e:
        logging.error(f"主程序执行失败: {e}")